import os

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm


def delete_files_in_folder(folder_path):
    files = os.listdir(folder_path)
    for file_name in files:
        file_path = os.path.join(folder_path, file_name)
        try:
            if os.path.isfile(file_path):
                os.remove(file_path)
            elif os.path.isdir(file_path):
                delete_files_in_folder(file_path)
        except Exception as e:
            print(f"Error deleting {file_path}: {e}")


delete_files_in_folder("tf-logs")

writer = SummaryWriter("tf-logs")
from data_set import stock_prediction_data_set
from model import TransformerClassifier

# 检测是否支持 CUDA
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# 数据长度
window_size = 30
# 创建自定义数据集实例

# 序列长度
batch_size = 2
# 模型和参数
num_layers = 5
# 数据维度
d_model = 16
num_head = 16
dim_feedforward = 512
ff_hidden_dim = 128
output_dim = 2
dropout = 0.1

train_data_set = stock_prediction_data_set(
    data_dir="train",
    data_file_limit=0.1,
    window=window_size + 1,
    data_column=['openIndex', 'closeIndex', 'highIndex', 'lowIndex',
                 'turnoverValueIndex', 'turnoverVolIndex',
                 'z2', 'z10', 'z30', 'm2', 'm10', 'm30',
                 'openHl', 'closeHl', 'highHl', 'lowHl'],
    label_column="profit", cache_label="train")

validate_data_set = stock_prediction_data_set(
    data_dir="validate",
    data_file_limit=0.1,
    window=window_size + 1,
    data_column=['openIndex', 'closeIndex', 'highIndex', 'lowIndex',
                 'turnoverValueIndex', 'turnoverVolIndex',
                 'z2', 'z10', 'z30', 'm2', 'm10', 'm30',
                 'openHl', 'closeHl', 'highHl', 'lowHl'],
    label_column="profit", cache_label="validate")
train_loader = DataLoader(train_data_set, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(validate_data_set, batch_size=batch_size, shuffle=False)

# 创建模型
model = TransformerClassifier(num_layers, d_model, window_size, num_head, dim_feedforward, ff_hidden_dim,
                              output_dim, device=device).to(device)

# 在日志中输出网络结构
writer.add_graph(model, torch.randn(batch_size, window_size, d_model))

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.00008)

# 训练模型
num_epochs = 1000
min_loss = 9999
train_idx = 0
for epoch in tqdm(range(num_epochs), desc="训练epoch", position=0, leave=True):
    model.train()  # 设置模型为训练模式
    total_loss = 0

    for batch_data, batch_labels in tqdm(train_loader, desc="训练batch", position=0, leave=True):
        batch_data = batch_data.to(device)
        batch_labels = batch_labels.to(device)
        optimizer.zero_grad()  # 清零梯度
        # 前向传播
        outputs = model(batch_data)
        # 计算损失
        loss = criterion(outputs, batch_labels)
        # 反向传播
        loss.backward()
        # 更新权重
        optimizer.step()
        # 损失相加
        loss_value = loss.item()
        total_loss += loss_value
        writer.add_scalar("each batch loss", loss_value, train_idx)
        train_idx += 1
    average_loss = total_loss / len(train_loader)
    writer.add_scalar("each epoch avg loss", average_loss, epoch)
    print(f"Epoch {epoch + 1}/{num_epochs}, Training Loss: {average_loss}")

    # 在验证集上评估模型
    model.eval()  # 设置模型为评估模式
    val_loss = 0
    correct = 0
    total = 0
    validate_idx = 0
    with torch.no_grad():
        for val_batch_data, val_batch_labels in tqdm(val_loader, desc="验证batch", position=0, leave=True):
            val_batch_labels = val_batch_labels.to(device)
            val_outputs = model(val_batch_data.to(device))
            val_loss_value = criterion(val_outputs, val_batch_labels).item()
            writer.add_scalar("each validate loss", val_loss_value, validate_idx)
            val_loss += val_loss_value
            # 比较第一列和第二列，根据条件设置新的值
            tensor_result = torch.zeros_like(val_outputs)
            tensor_result[:, 0] = torch.where(val_outputs[:, 0] > val_outputs[:, 1], 1, 0)
            tensor_result[:, 1] = 1 - tensor_result[:, 0]

            predicted = tensor_result.float().to(device)
            total += val_batch_labels.size(0)
            correct += ((predicted == val_batch_labels).sum().item() / 2)
            validate_idx += 1
    # 打印在验证集上的性能指标
    val_accuracy = correct / total
    val_loss_result = val_loss / len(val_loader)
    writer.add_scalar("each epoch validate loss", val_loss_result, epoch)
    writer.add_scalar("each epoch validate accuracy", val_accuracy, epoch)
    print(
        f"Epoch {epoch + 1}/{num_epochs}, Validation Loss: {val_loss_result}, Validation Accuracy: {val_accuracy}")

    if val_loss_result < min_loss:
        # 保存训练好的模型
        torch.save(model.state_dict(), "transformer_classifier.pth")
        min_loss = val_loss_result
